import cv2
import numpy as np
import os

def preprocess_image(image_path, output_dir, target_size=640):
    # 加载图像
    image = cv2.imread(image_path)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # 转换为RGB格式

    # 等比例缩放
    imgWidth = image.shape[1]
    imgHeight = image.shape[0]
    scale = min(target_size / imgWidth, target_size / imgHeight)
    resized_width = int(imgWidth * scale)
    resized_height = int(imgHeight * scale)
    image = cv2.resize(image, (resized_width, resized_height), interpolation=cv2.INTER_AREA)

    # 居中对齐并填充
    top, bottom = (target_size - resized_height) // 2, (target_size - resized_height + 1) // 2
    left, right = (target_size - resized_width) // 2, (target_size - resized_width + 1) // 2
    image = cv2.copyMakeBorder(image, top, bottom, left, right, cv2.BORDER_CONSTANT, value=[128, 128, 128])

    # 保存
    filename = os.path.basename(image_path)
    cv2.imwrite(output_dir + '/' + filename, image)

    # 归一化至0-1范围
    image = image.astype(np.float32) / 255.0

    # 调整维度顺序至C, H, W
    image = np.transpose(image, (2, 0, 1))

    # 加入批次维度
    image = np.expand_dims(image, axis=0)

    return image

def preprocess_images(image_dir, output_dir, target_size=640):
  # 获取文件夹内所有图片的名字
  images_names = os.listdir(image_dir)
  # 过滤出.jpg或.png图片（如果你有其他格式的图片，请相应修改）
  images_names = [img for img in images_names if img.endswith('.jpg') or img.endswith('.png')]
  images = []
  for img_name in images_names:
    image_path = os.path.join(image_dir, img_name)
    image = preprocess_image(image_path, output_dir, target_size)
    images.append(image)
  return images

if __name__ == "__main__":
  # 单个示例使用
  image_path = '../yolo-cow/datasets/cow/images/train/000001000001.jpg'
  output_dir = './output'
  processed_image = preprocess_image(image_path, output_dir)
  # print(processed_image.shape)  # 输出应为 (1, 3, 640, 640)

  # 多个示例使用
  # preprocess_images('../yolo-cow/datasets/cow/images/train', './output')
